'''
Descripttion: 基于PCA与EM算法的多光谱遥感影像变化检测研究_吴柯 论文复现
Author: Haixu He
Date: 2021-12-21 20:10:18
'''
import numpy as np
import matplotlib.pyplot as plt
import model.em as em
import math
import cmath


def BinaryPrimary(a, b, c):
    x1 = None
    x2 = None
    discriminant = (b**2) - (4 * a * c)
    if discriminant == 0:
        x1 = -(b / (2 * a))
    else:
        if discriminant > 0:
            root = math.sqrt(discriminant)
        else:
            root = cmath.sqrt(discriminant)
        x1 = (-b + root) / (2 * a)
        x2 = (-b - root) / (2 * a)
    equation = ("{0}x\N{SUPERSCRIPT TWO}+{1}x+{2}=0" " \N{RIGHTWARDS ARROW} x={3}").format(a, b, c, x1)
    # \N{RIGHTWARDS ARROW} 代表显示一个箭头标识(→)
    if x2 is not None:
        equation += ' or x={0}'.format(x2)
    return x1, x2


if __name__ == '__main__':
    alpha = 0.5
    data = np.abs(np.load('data/r.npy'))
    plt.plot(data)

    er_max, er_min = np.max(data), np.min(data)
    """公式7/8/9"""
    M_d = (er_max - er_min) / 2
    T_n = (1 - alpha) * M_d
    T_c = (1 + alpha) * M_d

    S_n = data[data < T_n]
    S_c = data[data >= T_c]
    """公式10"""
    p_init_n = S_n.shape[0] / data.shape[0]
    m_init_n = np.sum(S_n) / S_n.shape[0]
    sigma_init_n = np.sum((S_n - m_init_n)**2) / S_n.shape[0]

    p_init_c = S_c.shape[0] / data.shape[0]
    m_init_c = np.sum(S_c) / S_c.shape[0]
    sigma_init_c = np.sum((S_c - m_init_c)**2) / S_c.shape[0]

    args = [p_init_n, m_init_n, sigma_init_n, p_init_c, m_init_c, sigma_init_c]

    # args = [0.6, 0.1, 1, 0.4, 0.1, 1]

    # 开始EM算法，进行参数估计
    p_n, m_n, sigma_n, p_c, m_c, sigma_c = em.EM_Train(data, args, iter=500)
    """求解公式11"""
    a = sigma_n - sigma_c
    b = 2 * (m_n * sigma_c - m_c * sigma_n)
    c = sigma_n * m_c**2 - sigma_c * m_n**2 - 2 * sigma_n * sigma_c * math.log(
        (math.sqrt(sigma_n) * p_n) / (math.sqrt(sigma_c) * p_c))

    x1, x2 = BinaryPrimary(a, b, c)
    print("求解的T0为：{}和{}".format(x1, x2))
    plt.show()
